#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
简单版URPC2020数据集训练脚本
"""

import argparse
from ultralytics import YOLO


def main():
    parser = argparse.ArgumentParser(description='URPC2020 YOLO Training - Simple')
    parser.add_argument('--model', type=str, default='yolov8n.pt', 
                       help='Base YOLO model (default: yolov8n.pt)')
    parser.add_argument('--data', type=str, default='urpc2020.yaml', help='Dataset config')
    parser.add_argument('--epochs', type=int, default=120, help='Training epochs')
    parser.add_argument('--batch', type=int, default=16, help='Batch size')
    parser.add_argument('--imgsz', type=int, default=640, help='Image size')
    parser.add_argument('--device', type=str, default='0', help='Device (0, 1, 2, 3 or cpu)')
    parser.add_argument('--project', type=str, default='runs/train', help='Project name')
    parser.add_argument('--name', type=str, default='urpc2020_yolov8', help='Experiment name')
    parser.add_argument('--exist-ok', action='store_true', help='Overwrite existing project/name')
    
    args = parser.parse_args()
    
    print("=" * 70)
    print("🌊 URPC2020水下目标检测训练")
    print("=" * 70)
    print("")
    
    print("⚙️  训练参数:")
    print(f"   模型: {args.model}")
    print(f"   数据: {args.data}")
    print(f"   轮数: {args.epochs}")
    print(f"   批次: {args.batch}")
    print(f"   图像尺寸: {args.imgsz}px")
    print(f"   设备: {args.device}")
    print("")
    
    # 创建YOLO模型
    model = YOLO(args.model)
    
    # 开始训练
    results = model.train(
        data=args.data,
        epochs=args.epochs,
        batch=args.batch,
        imgsz=args.imgsz,
        device=args.device,
        project=args.project,
        name=args.name,
        exist_ok=args.exist_ok,
    )
    
    print("")
    print("=" * 70)
    print("✅ 训练完成！")
    print("=" * 70)
    print("")
    print("📝 查看结果:")
    print(f"   训练目录: {args.project}/{args.name}/")
    print(f"   最佳模型: {args.project}/{args.name}/weights/best.pt")
    print(f"   最终模型: {args.project}/{args.name}/weights/last.pt")
    print("")


if __name__ == '__main__':
    main()
